#include "force-unroll.hpp"
#include <expr.h>
#include <explow.h>
#include <gimple.h>
#include <gimple-pretty-print.h>
#include <map>
#include <print-tree.h>
#include <diagnostic.h>
#include <gimple-ssa.h>
#include <tree-vrp.h>
#include <stringpool.h>
#include <tree-ssanames.h>
#include <toplev.h>
#include <cfgloop.h>
#include <tree-ssa-loop-manip.h>
#include <tree-cfg.h>
#include <cfg.h>
#include <set>
typedef int loop_id;
loop_id get_loop_id(struct loop *loop){
  return loop->num;
}

namespace force_unroll{
  std::map<loop_id, unsigned int> loop_unroll_records;
  unsigned int loop_unroll_adjust(unsigned unroll, struct loop *loop){
    loop_id id = get_loop_id(loop);
    if (loop_unroll_records.find(id) != loop_unroll_records.end()){
      niter_desc *desc = get_simple_loop_desc(loop);
      if (!desc->simple_p || !desc->const_iter || desc->assumptions) {
        loop->lpt_decision.decision = LPT_UNROLL_CONSTANT;
        loop->lpt_decision.times = loop_unroll_records[id] - 1;
      }
      //printf("safelen, times: %d %d\n", loop->safelen, loop->simple_loop_desc->niter);
      
      /* 临时以UNROLL_STUPID决定，以避免调整后一些傻不拉几的判断 */
      /* 不行，UNROLL_STUPID太蠢了。。。唉*/
      // loop->lpt_decision.decision = LPT_UNROLL_STUPID;
      // loop->lpt_decision.times = loop_unroll_records[id] - 1;
      return loop_unroll_records[id];
    } else if (nohp::targetm_default.loop_unroll_adjust)
      return nohp::targetm_default.loop_unroll_adjust(unroll, loop);
    return unroll;
  }
}
static tree gdecl;
class builtin_force_unroll : public nohp_builtin_function{
public:
  builtin_force_unroll() {
    name = "__builtin_force_unroll";
  }
  tree type() override {
    return build_function_type_list(void_type_node, intSI_type_node, NULL_TREE);
  }
  tree fold(tree fndecl, int n_args, tree *argp, bool ignore) override{
    return NULL_TREE;
  }
  /*虽然gimple_fold不如expand的效果稳定，但是必须消掉这种builtin才能向量化，另一个思路是声明对应的向量builtin*/
  bool gimple_fold(gimple_stmt_iterator *gsi) override{
    gimple* stmt = gsi_stmt(*gsi);
    if (stmt->bb && stmt->bb->loop_father){
      auto loop = stmt->bb->loop_father;
      auto times = gimple_call_arg(stmt, 0);
      gcc_assert(TREE_CONSTANT(times));
      //printf("f: %p %p %d %s\n", loop->header, loop->former_header, loop->header->index, current_pass->name);
      loop_id id = get_loop_id(loop);
      force_unroll::loop_unroll_records[id] = TREE_INT_CST_LOW(times);//XINT(16, 0);
      /*关键：释放原stmt的def，不然一会依赖分析的时候就boom了*/
      unlink_stmt_vdef(stmt);
      release_defs(stmt);
      auto nop = gimple_build_nop();
      gimple_set_location(nop, gimple_location(stmt));
      gsi_replace(gsi, nop, true);
      return true;
    } else {
      return false;
    }
  }
  rtx expand(tree exp, rtx target,
             rtx subtarget,
             machine_mode mode,
             int ignore) override{
    call_expr_arg_iterator iter;
    auto op = expand_expr(first_call_expr_arg(exp, &iter), NULL_RTX, VOIDmode, EXPAND_CONST_ADDRESS);
    if (!CONST_INT_P(op))
      error("Unroll times must be an integer literal");
    auto insn = emit_move_insn(gen_rtx_REG(DImode, 31), gen_rtx_REG(DImode, 31));
    auto loop =  currently_expanding_gimple_stmt->bb->loop_father;

    printf("e: %p %p\n", loop->header, loop->former_header);
    loop_id id = get_loop_id(loop);
    force_unroll::loop_unroll_records[id] = XINT(op, 0);

    return NULL_RTX;
  }
  void post_decl() override {
    gdecl = decl;
  }
};

// namespace {
//   const pass_data pass_data_sw_64_force_unroll = {
//     GIMPLE_PASS,
//     "sw_64_force_unroll",
//     OPTGROUP_LOOP,
//     TV_NONE,
//     (PROP_cfg | PROP_ssa | PROP_loops),
//     0,
//     0,
//     0,
//     TODO_update_ssa
//   };
// };
// #include <tree-ssa-loop.h>
// #include <tree-ssa-loop-manip.h>
// #include <tree-ssa-loop-ivopts.h>
// class pass_sw_64_force_unroll : public gimple_opt_pass {
// public:
//   pass_sw_64_force_unroll(gcc::context *ctxt) : gimple_opt_pass(pass_data_sw_64_force_unroll, ctxt){}
//   virtual bool gate(function *f){
//     return true;
//   }
//   virtual unsigned int execute(function *f){
//     struct loop *loop;
//     FOR_EACH_LOOP(loop, 0){
//       int nbbs = loop->num_nodes;
//       basic_block *bbs = get_loop_body(loop);
//       int unroll = -1;
//       //printf("%p %d\n", loop, gsi_stmt(gsi_start_bb(loop->header))->location);

//       for (int i = 0; i < nbbs; i ++){
//         basic_block bb = bbs[i];
//         for (auto si = gsi_start_bb(bb); !gsi_end_p(si); gsi_next(&si)){
//           gimple *stmt = gsi_stmt(si);
//           if (gimple_code(stmt) == GIMPLE_CALL) {
//             tree fn = gimple_call_fndecl(stmt);

//             if (!strcmp(IDENTIFIER_POINTER(DECL_NAME(fn)), "__builtin_force_unroll")) {
//               unroll = TREE_INT_CST_LOW(gimple_call_arg(stmt, 0));
//               // unlink_stmt_vdef(stmt);
//               // release_defs(stmt);
//               // auto nop = gimple_build_nop();
//               // gimple_set_location(nop, gimple_location(stmt));
//               // gsi_replace(&si, nop, true);

//               //debug(gimple_call_arg(stmt, 0));
//             }
//           }
//         }
//       }
//       if (unroll != -1){
//         tree_niter_desc tidesc;
//         printf("%p\n", single_dom_exit(loop));
//         if (can_unroll_loop_p(loop, unroll, &tidesc)) {
//           puts("loop can unroll");
//           tree_unroll_loop(loop, unroll, single_dom_exit(loop), &tidesc);
//         } else {
//           printf("cannot unroll %d times\n", unroll);
//         }
//       }
//     }
//     return 0;
//   }
// };

#include <context.h>
void sw_64_force_unroll::add_builtins(){
  targetm.loop_unroll_adjust = force_unroll::loop_unroll_adjust;
  auto force_unroll = new builtin_force_unroll;
  nohp::add_builtin(force_unroll);
  //nohp::register_pass(new pass_sw_64_force_unroll(g), "ivopts", 1, PASS_POS_INSERT_AFTER);
}
